#!/usr/bin/env node

/**
 * MySQL MCP 服务器
 * 这是一个基于 Model Context Protocol (MCP) 的 MySQL 数据库连接服务
 * 
 * 主要功能：
 * - 连接 MySQL 数据库
 * - 执行 SQL 查询语句
 * - 返回查询结果
 * - 支持自定义数据库连接参数
 *
 * @author 田小檬 
 * @Email <admin@lmwa.cn>
 * @version 1.0
 * @date 2025/9/2 22:40:10
 */

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import {
  CallToolRequestSchema,
  ErrorCode,
  ListToolsRequestSchema,
  McpError,
} from '@modelcontextprotocol/sdk/types.js';
import mysql from 'mysql2/promise';

/**
 * MySQL 连接池管理类
 */
class MySQLManager {
  constructor() {
    this.pools = new Map();
  }

  /**
   * 获取或创建数据库连接池
   * @param {Object} config - 数据库配置
   * @returns {Object} 连接池实例
   */
  getPool(config) {
    const key = `${config.host}:${config.port}:${config.database}:${config.user}`;
    
    if (!this.pools.has(key)) {
      const pool = mysql.createPool({
        host: config.host || 'localhost',
        port: config.port || 3306,
        user: config.user || 'root',
        password: config.password || '',
        database: config.database,
        waitForConnections: true,
        connectionLimit: 10,
        queueLimit: 0,
        acquireTimeout: 60000,
        timeout: 60000
      });
      this.pools.set(key, pool);
    }
    
    return this.pools.get(key);
  }

  /**
   * 执行 SQL 查询
   * @param {Object} config - 数据库配置
   * @param {string} sql - SQL 语句
   * @returns {Object} 查询结果
   */
  async executeQuery(config, sql) {
    const pool = this.getPool(config);
    
    try {
      const [rows, fields] = await pool.execute(sql);
      return {
        success: true,
        data: rows,
        fields: fields ? fields.map(field => ({
          name: field.name,
          type: field.type,
          length: field.length
        })) : [],
        rowCount: Array.isArray(rows) ? rows.length : 0,
        affectedRows: rows.affectedRows || 0,
        insertId: rows.insertId || null
      };
    } catch (error) {
      return {
        success: false,
        error: error.message,
        code: error.code,
        errno: error.errno,
        sqlState: error.sqlState
      };
    }
  }

  /**
   * 关闭所有连接池
   */
  async closeAll() {
    for (const pool of this.pools.values()) {
      await pool.end();
    }
    this.pools.clear();
  }
}

// 创建 MySQL 管理器实例
const mysqlManager = new MySQLManager();

/**
 * 创建并配置 MCP 服务器
 */
const server = new Server(
  {
    name: 'xiaomeng-mysql-server',
    version: '1.0.0',
  },
  {
    capabilities: {
      tools: {},
    },
  }
);

/**
 * 注册工具列表处理器
 */
server.setRequestHandler(ListToolsRequestSchema, async () => {
  return {
    tools: [
      {
        name: 'execute_sql',
        description: '执行 MySQL SQL 语句并返回结果',
        inputSchema: {
          type: 'object',
          properties: {
            host: {
              type: 'string',
              description: 'MySQL 服务器地址（默认：localhost）',
              default: 'localhost'
            },
            port: {
              type: 'number',
              description: 'MySQL 服务器端口（默认：3306）',
              default: 3306
            },
            database: {
              type: 'string',
              description: '数据库名称'
            },
            user: {
              type: 'string',
              description: '用户名（默认：root）',
              default: 'root'
            },
            password: {
              type: 'string',
              description: '密码'
            },
            sql: {
              type: 'string',
              description: '要执行的 SQL 语句'
            }
          },
          required: ['database', 'sql'],
          additionalProperties: false
        }
      },
      {
        name: 'test_connection',
        description: '测试 MySQL 数据库连接',
        inputSchema: {
          type: 'object',
          properties: {
            host: {
              type: 'string',
              description: 'MySQL 服务器地址（默认：localhost）',
              default: 'localhost'
            },
            port: {
              type: 'number',
              description: 'MySQL 服务器端口（默认：3306）',
              default: 3306
            },
            database: {
              type: 'string',
              description: '数据库名称'
            },
            user: {
              type: 'string',
              description: '用户名（默认：root）',
              default: 'root'
            },
            password: {
              type: 'string',
              description: '密码'
            }
          },
          required: ['database'],
          additionalProperties: false
        }
      }
    ],
  };
});

/**
 * 注册工具调用处理器
 */
server.setRequestHandler(CallToolRequestSchema, async (request) => {
  const { name, arguments: args } = request.params;

  try {
    switch (name) {
      case 'execute_sql': {
        const config = {
          host: args.host || 'localhost',
          port: args.port || 3306,
          database: args.database,
          user: args.user || 'root',
          password: args.password || ''
        };

        const result = await mysqlManager.executeQuery(config, args.sql);
        
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(result, null, 2)
            }
          ]
        };
      }

      case 'test_connection': {
        const config = {
          host: args.host || 'localhost',
          port: args.port || 3306,
          database: args.database,
          user: args.user || 'root',
          password: args.password || ''
        };

        try {
          const pool = mysqlManager.getPool(config);
          const connection = await pool.getConnection();
          connection.release();
          
          return {
            content: [
              {
                type: 'text',
                text: JSON.stringify({
                  success: true,
                  message: '数据库连接测试成功',
                  config: {
                    host: config.host,
                    port: config.port,
                    database: config.database,
                    user: config.user
                  }
                }, null, 2)
              }
            ]
          };
        } catch (error) {
          return {
            content: [
              {
                type: 'text',
                text: JSON.stringify({
                  success: false,
                  message: '数据库连接测试失败',
                  error: error.message,
                  config: {
                    host: config.host,
                    port: config.port,
                    database: config.database,
                    user: config.user
                  }
                }, null, 2)
              }
            ]
          };
        }
      }

      default:
        throw new McpError(
          ErrorCode.MethodNotFound,
          `未知的工具: ${name}`
        );
    }
  } catch (error) {
    throw new McpError(
      ErrorCode.InternalError,
      `工具执行错误: ${error.message}`
    );
  }
});

/**
 * 启动服务器
 */
async function main() {
  const transport = new StdioServerTransport();
  await server.connect(transport);
  
  // 优雅关闭处理
  process.on('SIGINT', async () => {
    console.error('收到 SIGINT 信号，正在关闭服务器...');
    await mysqlManager.closeAll();
    process.exit(0);
  });
  
  process.on('SIGTERM', async () => {
    console.error('收到 SIGTERM 信号，正在关闭服务器...');
    await mysqlManager.closeAll();
    process.exit(0);
  });
}

// 启动服务器
main().catch((error) => {
  console.error('服务器启动失败:', error);
  process.exit(1);
});